from db.es.es_service import ElasticsearchStorage
import pandas as pd
from datetime import datetime
import settings

# 创建Elasticsearch客户端
es = ElasticsearchStorage()


def find_session_with_text(text_info):
    text_info = '我们会根据这些信息制定一个更加个性化的退休规划方案'
    query = {
        "query": {
            "match_phrase": {
                "text_info": {
                    "query": text_info
                }
            }
        }
    }
    results = es.search_documents('dialogues', query)
    df = pd.DataFrame(columns=["对话id", "角色", "对话内容", "对话时间", "态度", "踩的原因", "用户建议"])
    if results:
        session_id = ''
        # 提取查询结果
        for hit in results:
            session_id = hit["_source"]["session_id"]

        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "term": {
                                "session_id": {
                                    "value": session_id
                                }
                            }
                        }
                    ]
                }
            }, "sort": [
                {
                    "message_time": {
                        "order": "asc"
                    }
                }
            ]
            , "size": 1000
        }
        results = es.search_documents('dialogues', query)
        if results:
            # 遍历查询结果
            for hit in results:
                session_id = hit["_source"]["session_id"]
                role = hit["_source"]["role"]
                text_info = hit["_source"]["text_info"]
                message_time = hit["_source"]["message_time"]
                attitude = hit["_source"].get("attitude", None)
                badoption = hit["_source"].get("badoption", None)
                userPrefer = hit["_source"].get("userPrefer", None)
                df = df.append({"对话id": session_id, "角色": role, "对话内容": text_info, "对话时间": message_time, "态度": attitude,
                                "踩的原因": badoption,
                                "用户建议": userPrefer},
                               ignore_index=True)

        # 导出到Excel文件
        df.to_excel("output.xlsx", index=False)


def find_session_with_date():
    session_ids = set([])
    # today = datetime.date.today()

    # 设置查询时间范围：从今天早上0点到晚上12点
    # start_time = today.strftime("%Y-%m-%dT00:00:00")
    # end_time = today.strftime("%Y-%m-%dT23:59:59")

    start_time_str = "2023-07-19 00:00:00"
    start_time = datetime.strptime(start_time_str, "%Y-%m-%d %H:%M:%S")

    end_time_str = "2023-07-20 00:00:00"
    end_time = datetime.strptime(end_time_str, "%Y-%m-%d %H:%M:%S")

    # 定义查询条件xw
    query = {
        "query": {
            "bool": {
                "must": [
                    {"range": {"message_time": {"gte": start_time, "lte": end_time}}},
                    # {"term": {"attitude": "bad"}}
                ]
            }
        }
        , "sort": [
            {
                "message_time": {
                    "order": "asc"
                }
            }
        ], "size": 1000
    }

    # 执行查询
    results = es.search_documents('dialogues', query)

    if results:
        session_id = ''
        # 创建一个空的DataFrame
        df = pd.DataFrame(columns=["对话id", "角色", "对话内容", "对话时间", "态度", "踩的原因", "用户建议"])
        # 提取查询结果
        for hit in results:
            session_id = hit["_source"]["session_id"]
            if session_id in session_ids:
                continue
            query = {
                "query": {
                    "bool": {
                        "must": [
                            {
                                "term": {
                                    "session_id": {
                                        "value": session_id
                                    }
                                }
                            }
                        ]
                    }
                }, "sort": [
                    {
                        "message_time": {
                            "order": "asc"
                        }
                    }
                ]
                , "size": 1000
            }
            results = es.search_documents('dialogues', query)
            if results:
                # 遍历查询结果
                for hit in results:
                    session_id = hit["_source"].get("session_id", None)
                    role = hit["_source"].get("role", None)
                    text_info = hit["_source"].get("text_info", None)
                    message_time = hit["_source"]["message_time"]
                    attitude = hit["_source"].get("attitude", None)
                    badoption = hit["_source"].get("badoption", None)
                    userPrefer = hit["_source"].get("userPrefer", None)
                    df = df.append(
                        {"对话id": session_id, "角色": role, "对话内容": text_info, "对话时间": message_time, "态度": attitude,
                         "踩的原因": badoption,
                         "用户建议": userPrefer},
                        ignore_index=True)
            session_ids.add(session_id)

        # 导出到Excel文件
        df.to_excel("output.xlsx", index=False)


def find_user_session_with_date():
    session_ids = set([])

    start_time_str = "2023-07-26 00:00:00"
    start_time = datetime.strptime(start_time_str, "%Y-%m-%d %H:%M:%S")

    end_time_str = "2023-07-27 00:00:00"
    end_time = datetime.strptime(end_time_str, "%Y-%m-%d %H:%M:%S")

    # 定义查询条件
    user_session_query = {
        "query": {
            "bool": {
                "must": [
                    {"range": {"message_time": {"gte": start_time, "lte": end_time}}},
                ]
            }
        }
        , "sort": [
            {
                "user_name": {
                    "order": "asc"
                }
            },
            {
                "message_time": {
                    "order": "asc"
                }
            }
        ], "size": 1000
    }

    # 执行查询
    session_results = es.search_documents('usersessions', user_session_query)

    if session_results:
        # 创建一个空的DataFrame
        df = pd.DataFrame(columns=["用户账号", "对话id", "角色", "对话内容", "对话时间", "态度", "踩的原因", "用户建议"])
        # 提取查询结果
        for hit in session_results:
            session_id = hit["_source"]["session_id"]
            user_name = hit["_source"]["user_name"]
            if session_id in session_ids:
                continue
            query = {
                "query": {
                    "bool": {
                        "must": [
                            {
                                "term": {
                                    "session_id": {
                                        "value": session_id
                                    }
                                }
                            }
                        ]
                    }
                }, "sort": [
                    {
                        "message_time": {
                            "order": "asc"
                        }
                    }
                ]
                , "size": 1000
            }
            results = es.search_documents('dialogues', query)
            if results:
                # 遍历查询结果
                for hit in results:
                    session_id = hit["_source"].get("session_id", None)
                    role = hit["_source"].get("role", None)
                    text_info = hit["_source"].get("text_info", None)
                    message_time = hit["_source"]["message_time"]
                    attitude = hit["_source"].get("attitude", None)
                    badoption = hit["_source"].get("badoption", None)
                    userPrefer = hit["_source"].get("userPrefer", None)

                    df = df.append(
                        {"用户账号": user_name, "对话id": session_id, "角色": role, "对话内容": text_info, "对话时间": message_time,
                         "态度": attitude,
                         "踩的原因": badoption,
                         "用户建议": userPrefer},
                        ignore_index=True)
            session_ids.add(session_id)

        def add_empty_row(group):
            return pd.concat([group, pd.DataFrame([[''] * len(group.columns)], columns=group.columns)],
                             ignore_index=True)

        df = df.groupby('用户账号').apply(add_empty_row).reset_index(drop=True)

        # 导出到Excel文件
        df.to_excel("output.xlsx", index=False)


if __name__ == '__main__':
    find_user_session_with_date()
